################################################################################
# 00-functions.R
# Paper:    Bottom Up? Top Down? Determinants of Issue-Attention in State 
#           Politics.
# Journal:  Journal of Politics.
# Authors:  Andreu Casas, Oscar Stuhler, Julia Payson, Joshua A. Tucker, 
#           Richard Bonneau, and Jonathan Nagler.
# Purpose:  A set of functions used in many of the other scripts:
#           - `transform_var_irf`: This function takes the impulse responses 
#                   estimated by vars::irf() and generates a dataframe with
#                   coefficients and bounds for both one-time shocks and 
#                   structural (permanent) impulses.
#           - `make_plot_data`: This function takes the dataframe generated by 
#             transform_var_irf and extracts from this only the coefficients to 
#             be shown in a subsequent plot. It generates a dataframe ready 
#             for generating a figure with results.
################################################################################

transform_var_irf <- function(var_irfs = var_irfs, days = 15){
  
  #-----------------------------------------------------------------------------
  # Function:    transform_var_irf
  # Input:
  #             var_irfs - an object with impulse response coefficients 
  #                        estimated by vars::irf()
  #             days     - the number of days to simulate
  # Output: 
  #             A dataframe with the responses to both, a one time 10% shock and 
  #               a permanent 10% increase
  #-----------------------------------------------------------------------------
  # Generate a character vector with the models endogenous variables
  variables <- names(var_irfs$irf)
  
  # Generate a character vector with elements to pull from the irf object
  elements_to_pull <- c("irf", "Upper", "Lower")
  
  # Initialize and fill a dataset with the irf information
  irf_data <- NULL
  for (el in elements_to_pull) {
    new_irf_info <- var_irfs[el][[1]]
    for (cov in variables) {
      new_irf_var_data <- as.data.frame(new_irf_info[cov][[1]])
      # Take the inverse logit to transform the effects to percentage point changes
      new_irf_var_data_transf <- as.data.frame(
        sapply(1:ncol(new_irf_var_data), function(j) inv.logit(new_irf_var_data[,j]) - 0.5))
      colnames(new_irf_var_data_transf) <- colnames(new_irf_var_data)
      new_irf_var_data_long <- new_irf_var_data_transf %>%
        gather(out, value)
      new_irf_var_data_long$cov <- cov
      new_irf_var_data_long$day <- rep(1:nrow(new_irf_var_data), 
                                       length(unique(new_irf_var_data_long$cov)))
      new_irf_var_data_long$e_type <- el
      irf_data <- rbind(irf_data, new_irf_var_data_long)
    }
  }
  
  # Rename 
  irf_data$e_type <- recode(irf_data$e_type, `irf` = "pe", `Lower` = "lwr", `Upper` = "upr")
  
  # Initializing an output dataset that will include IRF coefficients for one-time and permanent shocks
  new_irf_data <- NULL
  
  # Get character vector with the group variables
  variables <- unique(irf_data$cov)
  
  irf_data <- irf_data %>%
    filter(day <= (days + 1))
  
  # - iterating through covariates
  for (covariate in variables) {
    # -iterating through outcomes
    for (outcome in variables) {
      # - skipping when covariate and response are the same
      if (covariate != outcome) {
        # - initializing a cummulative-shocks matrix for this scenario: two 3-dim 
        #     matrix, one matrix for the covariate and one matrix for the response,
        #     and one dimension for the point estimate and the two other dimensions
        #     for the lower and upper bounds of the estimate
        cov_mat <- array(0, dim = c(days, days, 3))
        out_mat <- array(0, dim = c(days, days, 3))
        
        # - pull the full N-day IRFs for the endogenous covariate
        cov_resp <- irf_data %>%
          filter(cov == covariate, out == covariate) %>%
          # - remove 1st row: it's part of the set-up (response at day 0)
          filter(day != 1) %>%
          mutate(day = day -1)
        
        # - pull the full N-day IRFs for the particular outcome variable
        out_resp <- irf_data %>%
          filter(cov == covariate, out == outcome) %>%
          # - remove 1st row: it's part of the set-up (response at day 0)
          filter(day != 1) %>%
          mutate(day = day -1)
        
        # - transforming the 15-day IRFs for the covariate and outcome to a wide
        #   3-column format (one column per estimate type: pe, lwr, upr)
        or_cov_resp <- cov_resp %>%
          dplyr::select(day, value, e_type) %>%
          spread(e_type, value) %>%
          dplyr::select(-day)
        
        or_out_resp <- out_resp %>%
          dplyr::select(day, value, e_type) %>%
          spread(e_type, value) %>%
          dplyr::select(-day)
        
        # - fill up the first rows of the scenario matrices with the original 
        #   1-day shock responses
        cov_mat[1,,1:3] <- or_cov_resp %>%
          as.matrix()
        out_mat[1,,1:3] <- or_out_resp %>%
          as.matrix()
        
        for (i in 2:days) {
          # - iterating through the rest of the 15 days, beyond the 1st one
          # - checking first how much attention the covariate group is predicted 
          #   to pay to the issue in day i-1
          cov_att_pe <- sum(cov_mat[,(i-1),2])
          
          # - calculating how big a new shock needs to be in order for the 
          #   covariate group to keep its attention to 100%
          cov_new_shock <- 1 - cov_att_pe
          
          # - re-scaling the original 100 percentage point shock to the new shock
          cov_new_resp <- or_cov_resp[1:(days-(i-1)),] * cov_new_shock
          out_new_resp <- or_out_resp[1:(days-(i-1)),] * cov_new_shock
          
          # - adding the response to this new shock to the scenario matrices
          cov_mat[i,i:days,1:3] <- cov_new_resp %>%
            as.matrix()
          out_mat[i,i:days,1:3] <- out_new_resp %>%
            as.matrix()
        }
        # - saving the output for this cov --> out 
        new_rows <- rbind(
          data.frame(
            cov = covariate,
            value = colSums(out_mat[,,1]),
            out = outcome,
            day = 1:days,
            e_type = "lwr",
            data_type = "structural"),
          data.frame(
            cov = covariate,
            value = colSums(out_mat[,,2]),
            out = outcome,
            day = 1:days,
            e_type = "pe",
            data_type = "structural"),
          data.frame(
            cov = covariate,
            value = colSums(out_mat[,,3]),
            out = outcome,
            day = 1:days,
            e_type = "upr",
            data_type = "structural")
        )
        new_irf_data <- rbind(new_irf_data, new_rows)
      }
    }
  }
  
  # Merging the permanent shock IRFs with the 1-time-shock IRFs
  irf_data$data_type <- "one_time_shock"
  irf_data <- irf_data %>%
    # Take out day 1 which is just pre-setting and re-label
    filter(day != 1) %>% 
    mutate(day = day-1)
  
  all_irf_data <- rbind(irf_data, new_irf_data)
  
  # Removing cases in which covariate and outcome are the same
  all_irf_data <- all_irf_data %>%
    filter(cov != out)
  
  # Simulate one-time and structural shocks of 10 instead of 100 percentage pts.
  # Present the results in 0-100 scale instead of 0-1
  all_irf_data_wide <- all_irf_data %>%
    mutate(value = (value / 10) * 100) %>%
    spread(e_type, value)
  
  # Relabel the data type
  all_irf_data_wide$data_type <- recode(
    all_irf_data_wide$data_type,
    `one_time_shock` =  "Effect of a one time 10 percentage point attention increase at day 0",
    `structural` = "Effect of a structural 10 percentage point attention increase at day 0")
  
  return(all_irf_data_wide)
}







make_plot_data = function(final_input = final_input, which_effects = "one time", 
                          party_d = F, column_group = c("state_legislators"), 
                          which_day = 15){
  
  #-----------------------------------------------------------------------------
  # Input:
  #             final_input    - an object with the impulse response coefficients
  #                              as generated by transform_var_irf
  #             which_effects  - either "one time" for showing effects of a one 
  #                              time 10% shock or "structural" for showing the 
  #                              effects of a permanent 10% increase
  #             column_group   - which groups to have in the columns, on out 
  #                              case this is "state_legislators"
  #             which_day      - which day to computer cumulative IRFs for, 
  #                              default is 15
  #             party_d        - indicator whether to differentiate by party 
  #                              (Model II) or not (Models I and III)
  # Output: 
  #             A dataframe for coefficients ready for plotting
  #-----------------------------------------------------------------------------
  
  # Show only the the effects for a certain day, specified in the which-day variable
  new_db <- final_input %>%
    filter(day == which_day,
           grepl(which_effects, data_type))
  
  # Make vector with group labels
  group_set_cov = sort(unique(new_db$out))
  
  # Initializa data to plot, then iterate through all rown and columns groups
  plot_db <- NULL
  for (colgroup in column_group) {
    for (rowgroup in group_set_cov) {
      if(colgroup != rowgroup){
        # Pull the row group response to the column group response and vice versa
        col_to_row_db <- new_db %>% 
          filter(cov == colgroup, out == rowgroup)
        row_to_col_db <- new_db %>%
          filter(cov == rowgroup, out == colgroup)
        
        # Combine the responses, adding some new variables, and adding it to the 
        # new dataset to plot
        new_row = rbind(
          data.frame(
            colgroup = colgroup,
            rowgroup = rowgroup,
            direction = "Row response to column",
            pe = col_to_row_db$pe,
            lwr = col_to_row_db$lwr,
            upr = col_to_row_db$upr
          ),
          data.frame(
            colgroup = colgroup,
            rowgroup = rowgroup,
            direction = "Column response to row",
            pe = row_to_col_db$pe,
            lwr = row_to_col_db$lwr,
            upr = row_to_col_db$upr
          )
        )
        plot_db <- rbind(plot_db, new_row)
      }
    }
  }
  
  if(!party_d){
    # Relabel the covariates
    plot_db$colgroup <- recode(plot_db$colgroup,
                               `national_legislators` = "Members of\nCongress",
                               `state_legislators` = "State Legislators",
                               `state_partisans` = "State Partisans",
                               `national_media` = "National Media",
                               `state_media` = "State Media",
                               `Trump` = "Trump",
                               `state_random_partisans` = "General Public")
    
    plot_db$rowgroup <- recode(plot_db$rowgroup,
                               `national_legislators` = "Members of\nCongress",
                               `state_legislators` = "State Legislators",
                               `state_partisans` = "State Partisans",
                               `national_media` = "National Media",
                               `state_media` = "State Media",
                               `Trump` = "Trump",
                               `state_random_partisans` = "General Public")
    
    # Reorder the covariate and outcome categories
    plot_db$colgroup <- factor(plot_db$colgroup,
                               levels = c("Members of\nCongress",
                                          "National Media",
                                          "Trump",
                                          "State Legislators",
                                          "General Public",
                                          "State Partisans",
                                          "State Media"))
    
    plot_db$rowgroup <- factor(plot_db$rowgroup,
                               levels = c("Members of\nCongress",
                                          "National Media",
                                          "Trump",
                                          "State Legislators",
                                          "General Public",
                                          "State Partisans",
                                          "State Media"))
  } else {
    # - relabel the covariates so they fit/look better in the plot
    plot_db$colgroup <- recode(plot_db$colgroup,
               `national_legislators_democrat` = "Democrats\nin Congress",
               `national_legislators_republican` = "Republicans\nin Congress",
               `state_legislators_democrat` = "Democrat\nState Legislators",
               `state_legislators_republican` = "Republican\nState Legislators",
               `state_partisans_democrat` = "Democrat\nState Partisans",
               `state_partisans_republican` = "Republican\nState Partisans",
               `national_media_` = "National Media",
               `state_media_` = "State Media",
               `Trump_` = "Trump",
               `state_random_partisans_` = "General Public"
               )
    
    plot_db$rowgroup <- recode(plot_db$rowgroup,
               `national_legislators_democrat` = "Democrats\nin Congress",
               `national_legislators_republican` = "Republicans\nin Congress",
               `state_legislators_democrat` = "Democrat\nState Legislators",
               `state_legislators_republican` = "Republican\nState Legislators",
               `state_partisans_democrat` = "Democrat\nState Partisans",
               `state_partisans_republican` = "Republican\nState Partisans",
               `national_media_` = "National Media",
               `state_media_` = "State Media",
               `Trump_` = "Trump",
               `state_random_partisans_` = "General Public")

    # - reordering the covariate and outcome categories
    plot_db$colgroup <- factor(plot_db$colgroup,
                               levels = c("Democrats\nin Congress",
                                          "Republicans\nin Congress",
                                          "National Media",
                                          "Trump",
                                          "Democrat\nState Legislators",
                                          "Republican\nState Legislators",
                                          "Democrat\nState Partisans",
                                          "Republican\nState Partisans",
                                          "General Public",
                                          "State Media"
                                          ))
    
    plot_db$rowgroup <- factor(plot_db$rowgroup,
                               levels = c("Democrats\nin Congress",
                                          "Republicans\nin Congress",
                                          "National Media",
                                          "Trump",
                                          "Democrat\nState Legislators",
                                          "Republican\nState Legislators",
                                          "Democrat\nState Partisans",
                                          "Republican\nState Partisans",
                                          "General Public",
                                          "State Media"))
  }
  
  # Specify the names of effect direction
  plot_db$direction <- recode(plot_db$direction,
                              `Row response to column` = "Row groups' response",
                              `Column response to row` = "State legislators' response")
  
  return(plot_db)
  
}
